#ifndef CODEGEN_HPP
#define CODEGEN_HPP

#include <map>
#include <vector>
#include <string>
#include <cstring>
using namespace std;

#include "visitor.hpp"
#include "opcode.hpp"

enum ValueType {
    VT_INT,
    VT_DOULE,
    VT_STRING,
    VT_CODE_OBJECT,
    VT_UNIT,
    VT_TRUE,
    VT_FALSE,
};

class CodeObject;

union ConstantValue {
    long long int_value;
    string* str_value;
    double dbl_value;
    CodeObject* co_value;
};

class Constant {
public:
    Constant(ValueType vt, int v) {
        const_type = vt;
        value.int_value = v;
    }

    Constant(ValueType vt, char* v) {
        const_type = vt;
        value.str_value = new string(v);
    }

    Constant(ValueType vt, CodeObject* co) {
        const_type = vt;
        value.co_value = co;
    }

    Constant(ValueType vt) {
        const_type = vt;
        value.int_value = 0;
    }

    ValueType const_type;
    ConstantValue value;  

    bool operator < (const Constant& c) const {
        if (const_type < c.const_type) {
            return true;
        }

        if (const_type > c.const_type) {
            return false;
        }

        if (const_type == VT_INT) {
            return value.int_value < c.value.int_value;
        }
        else if (const_type == VT_DOULE) {
            return value.dbl_value < c.value.dbl_value;
        }
        else if (const_type == VT_STRING) {
            return value.str_value->compare(*c.value.str_value) < 0;
        }
        else if (const_type == VT_CODE_OBJECT) {
            return value.co_value < c.value.co_value;
        }

        return false;
    }
};

/*
 * Assistant class for generating if.
 */
class Label {
private:
    bool            _resolved;
    unsigned short  _pos;
    vector<int>     _uses;

public:
    Label() : _resolved(false), _pos(0) {}

    bool is_resolved()      { return _resolved; }
    void set_resolved(bool r) { _resolved = r; }
    void set_pos(int pos)   { _pos = (unsigned short)pos; }
    unsigned short pos()    { return _pos; }
    void add_use(int pos)   { _uses.push_back(pos); }

    void resolve(vector<unsigned char>& codes);
};

class SymbolTable;
class CodeObject;
class TypeChecker;

class CodeGen : public Visitor {
private:
    map<Constant, int>      _const_map;
    vector<Constant>        _consts;
    vector<unsigned char>   _insts;

    TypeChecker*            _type_checker;
    SymbolTable*            _sym_table;

    int                     _offset;
    int                     _context;

    int                     _arg_count;
    int                     _nlocals;
    int                     _stack_size;
    int                     _flags;

    bool                    _is_func;

public:
    CodeGen(SymbolTable* symtable, bool is_func = false);
    ~CodeGen();

    vector<unsigned char>& insts() { return _insts; }

    void add_op(unsigned char op_code, unsigned short param = 0);
    void jump_abs(unsigned char op_code, Label& label);
    void jump_rel(unsigned char op_code, Label& label);

    int compare_op(int op);

    void bind(Label& label);

    void prelogue(Node* n);
    void set_arg_count(int n)   { _arg_count = n; }

    void visit(Node* n);

    virtual void visit(PrintNode* n);
    virtual void visit(ReturnNode* n);
    virtual void visit(BinaryOp* n);
    virtual void visit(LogicNotNode* n);
    virtual void visit(ListNode* n);
    virtual void visit(ConstInt* n);
    virtual void visit(ConstString* n);
    virtual void visit(ConstArray* n);
    virtual void visit(IfNode* n);
    virtual void visit(WhileNode* n);
    virtual void visit(CmpNode* n);
    virtual void visit(AssignNode* n);
    virtual void visit(VarNode* n);
    virtual void visit(VarDefNode* n);
    virtual void visit(ParamDefNode* n);
    virtual void visit(FuncDefNode* n);
    virtual void visit(LambdaDef* n);
    virtual void visit(CallNode* n);
    virtual void visit(SubscrNode* n);
    virtual void visit(BuiltinConst* n);

    virtual void visit(ArrowTypeNode* n);
    virtual void visit(ArrayTypeNode* n);

    virtual void visit(TypeVarNode* n);
    virtual void visit(TypeFuncDefNode* n);
    virtual void visit(TypeVarDefNode* n);
    virtual void visit(TypeCallNode* n);
    virtual void visit(TypeConstInt* n);

    CodeObject* make_code_object();
    unsigned short put_consts_if_absent(Constant& c);
};

class Disassembler {
private:
    map<int, string> bytecodes;
public:
    #define ADD_PAIR(op)    bytecodes[op] = #op;
    Disassembler() {
        ADD_PAIR(UNARY_NOT);
        ADD_PAIR(BINARY_ADD);
        ADD_PAIR(BINARY_SUB);
        ADD_PAIR(BINARY_MUL);
        ADD_PAIR(BINARY_SUBSCR);
        ADD_PAIR(BINARY_DIV);
        ADD_PAIR(BINARY_AND);
        ADD_PAIR(BINARY_OR);
        ADD_PAIR(BINARY_XOR);
        ADD_PAIR(PRINT_NEWLINE);
        ADD_PAIR(PRINT_ITEM);
        ADD_PAIR(LOAD_CONST);
        ADD_PAIR(JUMP_ABSOLUTE);
        ADD_PAIR(POP_JUMP_IF_FALSE);
        ADD_PAIR(JUMP_FORWARD);
        ADD_PAIR(COMPARE_OP);
        ADD_PAIR(SETUP_LOOP);
        ADD_PAIR(POP_BLOCK);
        ADD_PAIR(STORE_SUBSCR);
        ADD_PAIR(STORE_NAME);
        ADD_PAIR(LOAD_NAME);
        ADD_PAIR(MAKE_FUNCTION);
        ADD_PAIR(CALL_FUNCTION);
        ADD_PAIR(BUILD_LIST);
    }

    void print_detail(unsigned char op, unsigned short param);
    void dis(vector<unsigned char>& insts);
};

#endif
